package hbbft

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"math"
	"math/rand"
	"sync"
	"time"
)

type HoneyBadger struct {
	Config
	// 当前epoch
	epoch uint64
	// epoch -> ACS
	acsInstances map[uint64]*ACS

	// 存事务的buffer
	transactionBuffer *transactionBuffer

	lock sync.Mutex
	// epoch -> []Transaction，在epoch中提交的事务
	outputs map[uint64][]Transaction

	// 真实网络接口
	network Network
	
    // 密码学组件（抽象后仍保留现有字段以最小侵入变更）
    TPKEPK *TPKEPublicKey
    TPKESK *TPKEPrivateKey
    TBLSPK *TBLSPublicKey
    TBLSSK *TBLSPrivateKey
    // 注入式提供者（后续real实现可替代）
    cryptoProvider CryptoProvider
	CommonCoin *CommonCoin
	HoneyBadgerBlock *HoneyBadgerBlock
	
	// 延迟监控
	delayMonitor *DelayMonitor
	
	// 性能监控
	perfMetrics *PerformanceMetrics
    logger      *Logger
    // 当前epoch起始时间，用于计算延迟
    epochStart  time.Time
	
	// 解密管理器
	// epoch -> batch_index -> proposer_id -> decryption_share
	decryptShares map[uint64]map[int]map[uint64]*DecryptionShare
	// epoch -> ACS输出的加密批次列表（待解密）
	pendingDecrypt map[uint64][]*EncryptedBatch
	// epoch -> 已解密完成的交易列表
	decryptedTxs map[uint64][]Transaction
}

func NewHoneyBadger(cfg Config, network Network) *HoneyBadger {
    // 通过构建标签选择 DKG 提供者：默认 mock，realcrypto 下为真实实现
    provider, err := defaultDKG.Setup(cfg.N, cfg.F+1)
    if err != nil { panic(fmt.Sprintf("DKG setup failed: %v", err)) }

    // 从provider提取密钥（mock和real模式使用不同的提取逻辑）
    tpkepk, tpsk, tblspk, tblssk := extractCryptoKeys(provider, cfg.ID, cfg)

    commonCoin := NewCommonCoin(cfg, tblspk, tblssk)
    honeyBadgerBlock := NewHoneyBadgerBlock(cfg, tpkepk, tpsk)
    
    // 初始化性能监控和日志
    perfMetrics := NewPerformanceMetrics(cfg.ID, cfg.N, cfg.F, cfg.BatchSize)
    logger := NewLogger(cfg.ID, "honeybadger", DEBUG)

	return &HoneyBadger{
        Config:            cfg,
        acsInstances:      make(map[uint64]*ACS),
        transactionBuffer: newTransactionBuffer(),
        outputs:           make(map[uint64][]Transaction),
        network:           network,
        TPKEPK:            tpkepk,
        TPKESK:            tpsk,
        TBLSPK:            tblspk,
        TBLSSK:            tblssk,
        CommonCoin:        commonCoin,
        HoneyBadgerBlock:  honeyBadgerBlock,
        delayMonitor:      NewDelayMonitor(),
        cryptoProvider:    provider,
        perfMetrics:       perfMetrics,
        logger:            logger,
        decryptShares:     make(map[uint64]map[int]map[uint64]*DecryptionShare),
        pendingDecrypt:    make(map[uint64][]*EncryptedBatch),
        decryptedTxs:      make(map[uint64][]Transaction),
        epochStart:        time.Now(), // 初始化第一个epoch的起始时间
    }
}

// SetNetwork 设置网络接口（用于延迟设置）
func (hb *HoneyBadger) SetNetwork(network Network) {
	hb.network = network
}

// GetMessage 已废弃，真实网络通过回调处理消息
func (hb *HoneyBadger) GetMessage() []Message {
	return []Message{}
}

// 加到buffer中
func (hb *HoneyBadger) AddTransaction(tx Transaction) {
	hb.transactionBuffer.add(tx)
	hb.logger.Debug("Transaction added to buffer", map[string]interface{}{
		"tx_hash": fmt.Sprintf("%x", tx.Hash()[:8]),
	})
}

// 处理给定epoch下的ACSMessage
func (hb *HoneyBadger) HandleMessage(sid, epoch uint64, msg *ACSMessage) error {
	acs, ok := hb.acsInstances[epoch]
	if !ok {
		if epoch < hb.epoch {
			return nil
		}
		acs = NewACS(hb.Config, hb.delayMonitor, epoch)
		hb.acsInstances[epoch] = acs
	}
	if err := acs.HandleMessage(sid, msg); err != nil {
		return err
	}
	hb.addMessages(acs.messageList.getMessages())

	if hb.epoch == epoch {
		acs, ok := hb.acsInstances[hb.epoch]
		if !ok {
			return nil
		}

		outputs := acs.Output()
		hb.logger.Debug("ACS Output check", map[string]interface{}{
			"epoch": hb.epoch,
			"outputs": outputs,
		})
		if outputs == nil || len(outputs) == 0 {
			return nil
		}

		// 判断是否跳过加密模式
		if hb.Config.SkipProofVerify {
			// 跳过加密模式：直接解码明文交易
			transactionMap := make(map[string]Transaction)
			for _, output := range outputs {
				// 跳过空的output
				if len(output) == 0 {
					continue
				}
				transactions := make([]Transaction, 0)
				if err := gob.NewDecoder(bytes.NewReader(output)).Decode(&transactions); err != nil {
					hb.logger.Warn("Failed to decode transactions, skipping", map[string]interface{}{
						"error": err.Error(),
						"output_len": len(output),
					})
					continue
				}
				for _, tx := range transactions {
					transactionMap[string(tx.Hash())] = tx
				}
			}

			transactions := make([]Transaction, 0)
			for _, tx := range transactionMap {
				transactions = append(transactions, tx)
			}

			hb.transactionBuffer.delete(transactions)
			hb.outputs[hb.epoch] = transactions
			
			// 记录epoch完成性能指标（真实延迟）
			latency := time.Duration(0)
			if !hb.epochStart.IsZero() {
				latency = time.Since(hb.epochStart)
			}
			// 获取各组件延迟
			rbcDelay := hb.delayMonitor.RBCDelays[hb.epoch]
			abaDelay := hb.delayMonitor.ABADelays[hb.epoch]
			coinDelay := hb.delayMonitor.CommonCoinDelays[hb.epoch]
			tpkeDelay := hb.delayMonitor.TPKEDelays[hb.epoch]
			
			hb.logger.Info("ACS decided", map[string]interface{}{
				"epoch":        hb.epoch,
				"txs":          len(transactions),
				"batch_size":   hb.BatchSize,
				"latency_ms":   latency.Milliseconds(),
				"rbc_ms":       rbcDelay.Milliseconds(),
				"aba_ms":       abaDelay.Milliseconds(),
				"coin_ms":      coinDelay.Milliseconds(),
				"tpke_ms":      tpkeDelay.Milliseconds(),
			})
			hb.perfMetrics.RecordEpochCompletion(hb.epoch, len(transactions), latency)
			hb.logger.LogEpochCompletion(hb.epoch, len(transactions), latency)
			
			hb.epoch++
			// 立即设置下一epoch的起始时间
			hb.epochStart = time.Now()

			return hb.propose()
		} else {
			// 完整加密模式：需要解密
			// 解析加密批次
			encryptedBatches := make([]*EncryptedBatch, 0)
			for batchIdx, output := range outputs {
				// 跳过空的output
				if len(output) == 0 {
					continue
				}
				var encBatch EncryptedBatch
				if err := gob.NewDecoder(bytes.NewReader(output)).Decode(&encBatch); err != nil {
					hb.logger.Warn("Failed to decode encrypted batch, skipping", map[string]interface{}{
						"error": err.Error(),
						"output_len": len(output),
						"batch_idx": batchIdx,
					})
					continue
				}
				encBatch.BatchIndex = int(batchIdx)
				encryptedBatches = append(encryptedBatches, &encBatch)
			}
			
			// 保存待解密批次
			hb.pendingDecrypt[hb.epoch] = encryptedBatches
			hb.logger.Info("ACS decided, starting decryption", map[string]interface{}{
				"epoch": hb.epoch,
				"batches": len(encryptedBatches),
			})
			
			// 生成解密份额并广播
			return hb.generateAndBroadcastDecryptionShares(hb.epoch, encryptedBatches)
		}
	}

	// 移除没用的acs
	for i, acs := range hb.acsInstances {
		if i >= hb.epoch-1 {
			continue
		}
		for _, t := range acs.bbaInstances {
			close(t.closeCh)
		}
		for _, t := range acs.rbcInstances {
			close(t.closeCh)
		}
		close(acs.closeCh)
		delete(hb.acsInstances, i)
	}

	return nil
}

// 调用propose
func (hb *HoneyBadger) Start() error {
	// 启动定期propose的goroutine
	go hb.proposeLoop()
	return hb.propose()
}

// proposeLoop 定期尝试propose
func (hb *HoneyBadger) proposeLoop() {
	ticker := time.NewTicker(2 * time.Second)
	defer ticker.Stop()
	
	for {
		select {
		case <-ticker.C:
			if err := hb.propose(); err != nil {
				hb.logger.Error("Failed to propose", map[string]interface{}{
					"error": err.Error(),
				})
			}
		}
	}
}

// 返回每次epoch已经提交的交易
func (hb *HoneyBadger) Outputs() map[uint64][]Transaction {
	hb.lock.Lock()
	defer hb.lock.Unlock()

	out := hb.outputs
	hb.outputs = make(map[uint64][]Transaction)
	return out
}

// 在当前epoch提出一批
func (hb *HoneyBadger) propose() error {
	if hb.transactionBuffer.len() == 0 {
		// No artificial delay - wait for real transactions
		hb.logger.Debug("No transactions in buffer, waiting...")
		return nil
	}
	
	hb.logger.Info("Starting new epoch proposal", map[string]interface{}{
		"epoch": hb.epoch,
		"buffer_size": hb.transactionBuffer.len(),
	})
	// 每个节点提议batchSize笔交易（不是batchSize/N）
	// 限制为buffer中实际可用的交易数量
	batchSize := hb.BatchSize
	availableTxs := hb.transactionBuffer.len()
	actualBatchSize := int(math.Min(float64(batchSize), float64(availableTxs)))
	
	// 确保至少有1笔交易
	if actualBatchSize == 0 {
		hb.logger.Warn("No transactions available for batch", nil)
		return nil
	}

	// 从buffer中随机选取actualBatchSize笔交易
	rand.Seed(time.Now().UnixNano())
	batch := make([]Transaction, actualBatchSize)
	for i := 0; i < actualBatchSize; i++ {
		// 从前availableTxs笔交易中随机选择
		idx := rand.Intn(availableTxs)
		batch[i] = hb.transactionBuffer.data[idx]
	}
	
	hb.logger.Info("Proposing batch", map[string]interface{}{
		"epoch": hb.epoch,
		"batch_size": actualBatchSize,
		"requested": batchSize,
		"available": availableTxs,
	})

	// 论文实验模式：跳过加密，直接序列化明文交易
	buf := new(bytes.Buffer)
	if hb.Config.SkipProofVerify {
		// 直接序列化交易列表
		if err := gob.NewEncoder(buf).Encode(batch); err != nil {
			return err
		}
	} else {
		// 完整加密模式：构建EncryptedBatch
		// 1. 序列化交易列表
		txsBuf := new(bytes.Buffer)
		if err := gob.NewEncoder(txsBuf).Encode(batch); err != nil {
			return fmt.Errorf("failed to encode transactions: %v", err)
		}
		txsData := txsBuf.Bytes()
		
		// 2. 生成随机AES密钥
		aesKey := randBytes(32)
		
		// 3. 使用AES加密交易数据
		cipherData, err := AESEncrypt(aesKey, txsData)
		if err != nil {
			return fmt.Errorf("failed to AES encrypt transactions: %v", err)
		}
		
		// 4. 使用TPKE加密AES密钥
		tkeyCiphertext, err := hb.TPKEPK.Encrypt(aesKey)
		if err != nil {
			return fmt.Errorf("failed to TPKE encrypt key: %v", err)
		}
		
		// 5. 构建EncryptedBatch（序列化密文组分）
		encBatch := EncryptedBatch{
			ProposerID: hb.ID,
			BatchIndex: 0, // 将由ACS完成后设置
			TKeyU:      serializeG1(tkeyCiphertext.U),
			TKeyV:      tkeyCiphertext.V,
			TKeyW:      serializeG2(tkeyCiphertext.W),
			CipherData: cipherData,
		}
		
		// 6. 序列化EncryptedBatch
		if err := gob.NewEncoder(buf).Encode(encBatch); err != nil {
			return fmt.Errorf("failed to encode EncryptedBatch: %v", err)
		}
	}

	acs, ok := hb.acsInstances[hb.epoch]
	if !ok {
		acs = NewACS(hb.Config, hb.delayMonitor, hb.epoch)
	}
	hb.acsInstances[hb.epoch] = acs

	if err := acs.InputValue(buf.Bytes()); err != nil {
		hb.logger.Error("Failed to input value to ACS", map[string]interface{}{
			"error": err.Error(),
		})
		return err
	}
	
	msgs := acs.messageList.getMessages()
	hb.logger.Debug("ACS generated messages", map[string]interface{}{
		"message_count": len(msgs),
	})
	
	hb.addMessages(msgs)
	return nil
}

func (hb *HoneyBadger) addMessages(msgs []Message) {
	for _, msg := range msgs {
		// 序列化消息
		hbMsg := HBMessage{hb.epoch, msg.Payload}
		msgBytes, err := hb.serializeMessage(hbMsg)
		if err != nil {
			hb.logger.Error("Failed to serialize message", map[string]interface{}{
				"error": err.Error(),
			})
			continue // 跳过序列化失败的消息
		}
		
		// 使用真实网络发送
        if msg.To == hb.ID {
			// 发送给自己，直接处理
            hb.handleIncomingMessage(hb.ID, msgBytes)
		} else {
			// 发送给其他节点
			hb.network.Send(msg.To, msgBytes)
			hb.perfMetrics.RecordNetworkMessage(true, len(msgBytes))
			hb.logger.LogNetworkMessage("sent", msg.To, len(msgBytes))
		}
	}
}

// serializeMessage 序列化 HoneyBadger 消息
func (hb *HoneyBadger) serializeMessage(msg HBMessage) ([]byte, error) {
	buf := new(bytes.Buffer)
	encoder := gob.NewEncoder(buf)
	if err := encoder.Encode(msg); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}

// handleIncomingMessage 处理接收到的消息
func (hb *HoneyBadger) handleIncomingMessage(senderID uint64, msgBytes []byte) error {
	var hbMsg HBMessage
	buf := bytes.NewReader(msgBytes)
	decoder := gob.NewDecoder(buf)
	if err := decoder.Decode(&hbMsg); err != nil {
		hb.logger.Error("Failed to deserialize message", map[string]interface{}{
			"error": err.Error(),
		})
		return err
	}
	
	// 记录接收到的网络消息
	hb.perfMetrics.RecordNetworkMessage(false, len(msgBytes))
    hb.logger.LogNetworkMessage("received", senderID, len(msgBytes))
	
	// 处理 ACS 消息
    if acsMsg, ok := hbMsg.Payload.(*ACSMessage); ok {
        return hb.HandleMessage(senderID, hbMsg.Epoch, acsMsg)
	}
	
	// 处理解密份额消息
	if decryptMsg, ok := hbMsg.Payload.(*DecryptShareMsg); ok {
		return hb.handleDecryptShareMsg(senderID, decryptMsg)
	}
	
	return nil
}

// GetPerformanceStats 获取性能统计信息
func (hb *HoneyBadger) GetPerformanceStats() map[string]interface{} {
	return hb.perfMetrics.GetStats()
}

// LogPerformanceStats 输出性能统计信息
func (hb *HoneyBadger) LogPerformanceStats() {
	hb.perfMetrics.LogStats()
}

// Close 关闭HoneyBadger并输出最终统计
func (hb *HoneyBadger) Close() {
	// 输出最终性能统计
	hb.LogPerformanceStats()
	
	// 关闭日志器
	if hb.logger != nil {
		hb.logger.Close()
	}
}

// generateAndBroadcastDecryptionShares 为加密批次生成解密份额并广播
func (hb *HoneyBadger) generateAndBroadcastDecryptionShares(epoch uint64, encryptedBatches []*EncryptedBatch) error {
	hb.logger.Info("Generating decryption shares", map[string]interface{}{
		"epoch": epoch,
		"batches": len(encryptedBatches),
	})
	
	// 为每个加密批次生成解密份额
	for _, encBatch := range encryptedBatches {
		if encBatch == nil {
			continue
		}
		
		// 反序列化TPKE密文
		ciphertext := &Ciphertext{
			U: deserializeG1(encBatch.TKeyU),
			V: encBatch.TKeyV,
			W: deserializeG2(encBatch.TKeyW),
		}
		
		// 生成解密份额
		share := hb.TPKESK.DecryptShare(ciphertext)
		
		// 序列化解密份额
		shareBytes := serializeG1(share.U_i)
		
		// 创建解密份额消息
		decryptMsg := &DecryptShareMsg{
			Epoch:      epoch,
			ProposerID: encBatch.ProposerID,
			BatchIndex: encBatch.BatchIndex,
			ShareU:     shareBytes,
		}
		
		// 广播解密份额给所有节点
		hbMsg := HBMessage{
			Epoch:   epoch,
			Payload: decryptMsg,
		}
		
		msgBytes, err := hb.serializeMessage(hbMsg)
		if err != nil {
			hb.logger.Error("Failed to serialize decrypt share message", map[string]interface{}{
				"error": err.Error(),
			})
			continue
		}
		
		// 广播给所有节点（包括自己）
		for _, nodeID := range hb.Nodes {
			if err := hb.network.Send(nodeID, msgBytes); err != nil {
				hb.logger.Error("Failed to send decrypt share", map[string]interface{}{
					"to": nodeID,
					"error": err.Error(),
				})
			}
		}
		
		hb.logger.Debug("Broadcasted decryption share", map[string]interface{}{
			"epoch": epoch,
			"proposer": encBatch.ProposerID,
			"batch_index": encBatch.BatchIndex,
		})
	}
	
	return nil
}

// handleDecryptShareMsg 处理接收到的解密份额消息
func (hb *HoneyBadger) handleDecryptShareMsg(senderID uint64, msg *DecryptShareMsg) error {
	hb.lock.Lock()
	defer hb.lock.Unlock()
	
	hb.logger.Debug("Received decryption share", map[string]interface{}{
		"from": senderID,
		"epoch": msg.Epoch,
		"proposer": msg.ProposerID,
		"batch_index": msg.BatchIndex,
	})
	
	// 初始化存储结构
	if hb.decryptShares[msg.Epoch] == nil {
		hb.decryptShares[msg.Epoch] = make(map[int]map[uint64]*DecryptionShare)
	}
	if hb.decryptShares[msg.Epoch][msg.BatchIndex] == nil {
		hb.decryptShares[msg.Epoch][msg.BatchIndex] = make(map[uint64]*DecryptionShare)
	}
	
	// 反序列化解密份额
	share := &DecryptionShare{
		U_i: deserializeG1(msg.ShareU),
	}
	
	// 存储解密份额
	hb.decryptShares[msg.Epoch][msg.BatchIndex][senderID] = share
	
	// 检查是否收集到足够的解密份额（需要 f+1 个）
	if len(hb.decryptShares[msg.Epoch][msg.BatchIndex]) >= hb.F+1 {
		hb.logger.Info("Collected enough decryption shares", map[string]interface{}{
			"epoch": msg.Epoch,
			"batch_index": msg.BatchIndex,
			"shares_count": len(hb.decryptShares[msg.Epoch][msg.BatchIndex]),
		})
		
		// 尝试解密该epoch的所有批次
		return hb.tryDecryptEpoch(msg.Epoch)
	}
	
	return nil
}

// tryDecryptEpoch 尝试解密一个epoch的所有批次
func (hb *HoneyBadger) tryDecryptEpoch(epoch uint64) error {
	// 检查该epoch是否有待解密的批次
	encryptedBatches, exists := hb.pendingDecrypt[epoch]
	if !exists || len(encryptedBatches) == 0 {
		return nil
	}
	
	// 检查是否所有批次都收集到了足够的解密份额
	for _, encBatch := range encryptedBatches {
		if len(hb.decryptShares[epoch][encBatch.BatchIndex]) < hb.F+1 {
			// 还有批次未收集够解密份额
			return nil
		}
	}
	
	hb.logger.Info("All batches ready for decryption", map[string]interface{}{
		"epoch": epoch,
		"batches": len(encryptedBatches),
	})
	
	// 所有批次都准备好了，开始解密
	allTransactions := make([]Transaction, 0)
	
	for _, encBatch := range encryptedBatches {
		// 重建密文
		ciphertext := &Ciphertext{
			U: deserializeG1(encBatch.TKeyU),
			V: encBatch.TKeyV,
			W: deserializeG2(encBatch.TKeyW),
		}
		
		// 获取该批次的解密份额
		shares := hb.decryptShares[epoch][encBatch.BatchIndex]
		
		// 转换为map[int]*DecryptionShare格式（CombineShares需要）
		sharesMap := make(map[int]*DecryptionShare)
		for nodeID, share := range shares {
			sharesMap[int(nodeID)] = share
		}
		
		// 组合解密份额恢复对称密钥
		key, err := hb.TPKEPK.CombineShares(ciphertext, sharesMap)
		if err != nil {
			hb.logger.Error("Failed to combine decryption shares", map[string]interface{}{
				"error": err.Error(),
				"batch_index": encBatch.BatchIndex,
			})
			continue
		}
		
		// 使用AES解密交易数据
		txsData, err := AESDecrypt(key, encBatch.CipherData)
		if err != nil {
			hb.logger.Error("Failed to decrypt batch data", map[string]interface{}{
				"error": err.Error(),
				"batch_index": encBatch.BatchIndex,
			})
			continue
		}
		
		// 反序列化交易列表
		var transactions []Transaction
		if err := gob.NewDecoder(bytes.NewReader(txsData)).Decode(&transactions); err != nil {
			hb.logger.Error("Failed to decode transactions", map[string]interface{}{
				"error": err.Error(),
				"batch_index": encBatch.BatchIndex,
			})
			continue
		}
		
		allTransactions = append(allTransactions, transactions...)
	}
	
	// 存储解密后的交易
	hb.outputs[epoch] = allTransactions
	
	// 记录epoch完成性能指标
	latency := time.Duration(0)
	if !hb.epochStart.IsZero() {
		latency = time.Since(hb.epochStart)
	}
	
	hb.logger.Info("Epoch decrypted and committed", map[string]interface{}{
		"epoch": epoch,
		"transactions": len(allTransactions),
		"latency_ms": latency.Milliseconds(),
	})
	
	hb.perfMetrics.RecordEpochCompletion(epoch, len(allTransactions), latency)
	hb.logger.LogEpochCompletion(epoch, len(allTransactions), latency)
	
	// 清理已完成epoch的解密数据
	delete(hb.pendingDecrypt, epoch)
	delete(hb.decryptShares, epoch)
	
	// 进入下一个epoch
	hb.epoch++
	// 立即设置下一epoch的起始时间
	hb.epochStart = time.Now()
	
	// 提议下一个epoch
	return hb.propose()
}
